#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName :start_browser_and_mcp_with_params.py
# @Created  :2025/04/05
# @Author   :批量启动Chromium+MCP实例（先启浏览器，后启MCP）

import os
import sys
import subprocess
import time
import socket
import threading
import argparse
import shutil
from pathlib import Path
from playwright.sync_api import sync_playwright

# 导入应用配置
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
from app.config import MCP_SERVERS_COUNT, PROJECT_PATH, DATA_PATH

# TCP端口转发相关常量和函数
BUFSIZE = 32 * 1024  # 单次最大缓冲 32 KiB


def log(msg):
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}")


def relay(sock_in: socket.socket, sock_out: socket.socket):
    """单向转发数据，直到 EOF"""
    try:
        while True:
            data = sock_in.recv(BUFSIZE)
            if not data:  # 对端关闭
                break
            sock_out.sendall(data)
    except OSError:
        pass
    finally:
        sock_in.close()
        sock_out.close()


def handle_client(client_conn: socket.socket, target_host: str, target_port: int):
    """接受一个客户端连接，并把它桥接到目标端口"""
    client_conn.settimeout(60)
    try:
        remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        client_conn.settimeout(60)
        remote.connect((target_host, target_port))
    except Exception as e:
        log(f"无法连接到目标 {target_host}:{target_port} -> {e}")
        client_conn.close()
        return

    # 启动两条线程，双向转发
    threading.Thread(target=relay, args=(client_conn, remote), daemon=True).start()
    threading.Thread(target=relay, args=(remote, client_conn), daemon=True).start()


def start_port_forwarding(listen_port, target_host, target_port):
    """启动端口转发服务"""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        server.bind(("", listen_port))
        server.listen(100)
        log(f"启动端口转发: *:{listen_port}  ->  {target_host}:{target_port}")

        # 作为守护线程运行，不阻塞主程序
        def forward_thread():
            try:
                while True:
                    conn, addr = server.accept()
                    log(f"新连接 {addr}")
                    threading.Thread(
                        target=handle_client,
                        args=(conn, target_host, target_port),
                        daemon=True
                    ).start()
            except Exception as e:
                log(f"端口转发服务异常: {e}")
            finally:
                server.close()

        threading.Thread(target=forward_thread, daemon=True).start()
        return True
    except OSError as e:
        log(f"无法启动端口转发服务 {listen_port}: {e}")
        return False


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='批量启动Chromium + MCP实例（一一对应）')
    parser.add_argument('--count', type=int, default=MCP_SERVERS_COUNT,
                        help=f'要启动的实例数量 (默认: {MCP_SERVERS_COUNT}，来自配置文件)')
    parser.add_argument('--mcp-start-port', type=int, default=8931,
                        help='MCP 起始端口号 (默认: 8931)')
    parser.add_argument('--cdp-start-port', type=int, default=9222,
                        help='Chromium CDP 调试端口起始号 (默认: 9222)')
    parser.add_argument('--start-user', type=int, default=1,
                        help='起始用户标识号 (默认: 1)')
    parser.add_argument('--host', type=str, default='127.0.0.1',
                        help='主机地址 (默认: 127.0.0.1)')
    parser.add_argument('--mcp-version', type=str, default='@playwright/mcp@0.0.44',
                        help='MCP版本 (默认: @playwright/mcp@0.0.44)')
    parser.add_argument('--delay', type=float, default=1.0,
                        help='启动实例之间的延迟(秒) (默认: 1.0)')
    parser.add_argument('--headless', action='store_true',
                        help='启用无头模式（Chromium 不显示窗口）')
    parser.add_argument('--is_start_mcp', action='store_true',
                        help='是否启动mcp')
    parser.add_argument('--is_start_cdp', action='store_true',
                        help='是否启动Chromium')
    parser.add_argument('--forward-port', type=int, default=19893,
                        help='端口转发的监听端口 (默认: 19893)')
    parser.add_argument('--forward-target-port', type=int, default=9222,
                        help='端口转发的目标端口 (默认: 9222)')
    parser.add_argument('--forward-target-host', type=str, default='127.0.0.1',
                        help='端口转发的目标主机 (默认: 127.0.0.1)')
    parser.add_argument('--enable-forward', action='store_true',
                        help='是否启用端口转发')
    return parser.parse_args()


def check_npx_availability():
    """检查npx命令是否可用（使用更可靠的方法，模仿CMD执行）"""
    if os.name == 'nt':  # Windows系统
        try:
            result = subprocess.run(
                ['cmd.exe', '/c', 'npx --version'],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=10
            )
            if result.returncode == 0:
                version = result.stdout.strip()
                print(f"✅ 已找到 npx 命令 (版本: {version})")
                return True
            else:
                print("⚠️  找不到 npx 命令。请确保已安装 Node.js 并添加到系统 PATH 中。")
                print("👉 下载地址: https://nodejs.org")
                return False
        except Exception as e:
            print(f"❌ 检查 npx 失败: {e}")
            return False
    else:  # 非Windows系统
        try:
            subprocess.run(['npx', '--version'], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            print("✅ 已找到 npx 命令")
            return True
        except (subprocess.SubprocessError, FileNotFoundError):
            print("⚠️  找不到 npx 命令，请确保 Node.js 已安装。")
            print("👉 下载地址: https://nodejs.org")
            return False


def get_chromium_executable_path():
    """使用 Playwright 获取 Chromium 可执行路径"""
    try:
        with sync_playwright() as p:
            exe_path = p.chromium.executable_path
            if not exe_path:
                print("❌ Playwright 未返回 Chromium 路径，请运行: npx playwright install chromium")
                return None
            exe = Path(exe_path)
            if not exe.exists():
                print(f"❌ Chromium 可执行文件不存在: {exe_path}")
                return None
            print(f"✅ 成功获取 Chromium 路径: {exe_path}")
            return str(exe_path)
    except Exception as e:
        print(f"❌ 获取 Chromium 路径失败: {e}")
        return None


def get_screen_size():
    """获取屏幕尺寸"""
    try:
        import tkinter as tk
        root = tk.Tk()
        root.withdraw()
        width = root.winfo_screenwidth()
        height = root.winfo_screenheight()
        root.destroy()
        return width, height
    except Exception:
        return 1920, 1080


def kill_existing_browser_processes():
    """
    关闭已存在的浏览器进程
    
    返回:
        int: 关闭的进程数量
    """
    try:
        import psutil
        killed_count = 0
        
        # 要查找的浏览器进程名
        browser_process_names = [
            'chrome.exe', 'chromium.exe', 'chrome', 'chromium',
            'google-chrome', 'google-chrome-stable'
        ]
        
        print("🔍 正在查找并关闭已存在的浏览器进程...")
        
        for proc in psutil.process_iter(['pid', 'name']):
            try:
                proc_name = proc.info['name'].lower()
                
                # 检查是否是浏览器进程
                for browser_name in browser_process_names:
                    if browser_name in proc_name:
                        print(f"🛑 发现浏览器进程: {proc_name} (PID: {proc.info['pid']})")
                        
                        # 终止进程
                        proc.kill()
                        proc.wait(timeout=5)  # 等待进程终止
                        print(f"✅ 已关闭进程: {proc_name} (PID: {proc.info['pid']})")
                        killed_count += 1
                        break
                        
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                # 进程可能已经结束，跳过
                continue
            except Exception as e:
                print(f"⚠️  处理进程时出错: {e}")
                continue
        
        if killed_count == 0:
            print("✅ 未发现需要关闭的浏览器进程")
        else:
            print(f"✅ 共关闭了 {killed_count} 个浏览器进程")
        
        return killed_count
        
    except ImportError:
        print("⚠️  未安装psutil库，无法自动关闭浏览器进程")
        print("💡 请手动关闭所有浏览器窗口，或运行: pip install psutil")
        return 0
    except Exception as e:
        print(f"❌ 关闭浏览器进程时出错: {e}")
        return 0


def start_chromium_instance(cdp_port, user_data_dir, headless=False):
    """启动单个 Chromium 实例"""
    # 直接使用指定的Chrome路径，不再通过Playwright获取
    chrome_path = "/opt/google/chrome/chrome"
    print(f"🔧 使用指定的Chrome路径: {chrome_path}")

    # 检查Chrome可执行文件是否存在
    # Chrome路径查找策略：
    # 1. 首先检查指定的默认路径
    # 2. 尝试从系统PATH中查找chrome或chromium
    # 3. 最后尝试通过playwright获取chromium路径
    if not os.path.exists(chrome_path):
        print(f"🔍 Chrome默认路径检查: {chrome_path} (不存在)")
        print(f"   提示: 这是{os.name == 'posix' and 'Linux' or 'Windows'}系统上Chrome的可能路径")
        print(
            f"   在{os.name == 'posix' and 'Linux' or 'Windows'}系统中，即使脚本不在Chrome目录下，也可以通过绝对路径启动Chrome")
        print(f"   正在尝试查找系统中可用的Chrome浏览器...")

        # 策略1: 尝试从系统PATH中查找chrome或chromium
        alt_chrome = shutil.which('chrome') or shutil.which('chromium')
        if alt_chrome:
            print(f"✅ 找到系统PATH中的Chrome/Chromium: {alt_chrome}")
            chrome_path = alt_chrome
        else:
            # 策略2: 尝试通过playwright获取chromium路径
            print("📦 尝试通过Playwright获取Chromium路径...")
            playwright_chrome = get_chromium_executable_path()
            if playwright_chrome:
                print(f"✅ 找到Playwright的Chromium: {playwright_chrome}")
                chrome_path = playwright_chrome
            else:
                # 所有策略都失败了
                error_msg = "❌ 无法找到Chrome或Chromium浏览器！\n"
                error_msg += f"   1. 确保已安装Chrome浏览器\n"
                error_msg += f"   2. 在Linux上，Chrome通常安装在 /opt/google/chrome/ 目录\n"
                error_msg += f"   3. 在Windows上，Chrome通常安装在 C:\\Program Files\\Google\\Chrome\\Application\\chrome.exe\n"
                error_msg += f"   4. 或者修改脚本中的chrome_path变量指向您的Chrome可执行文件"
                print(error_msg)
                raise RuntimeError("无法找到Chrome浏览器。请确保已安装Chrome并检查路径。")

    print(f"✅ 最终使用的Chrome路径: {chrome_path}")

    # 确保user_data_dir目录存在
    os.makedirs(user_data_dir, exist_ok=True)
    print(f"📁 浏览器使用用户数据目录: {user_data_dir}")

    # 使用用户指定的启动参数
    cmd = [
        chrome_path,
        f"--remote-debugging-port={cdp_port}",
        "--remote-debugging-address=0.0.0.0",
        f"--user-data-dir={user_data_dir}",
        "--no-sandbox",
        "--headless",  # 直接使用--headless参数
        "--disable-gpu",
        "--disable-dev-shm-usage",
        "--disable-software-rasterizer",
        "--ignore-certificate-errors",
        "--allow-running-insecure-content",
        "--disable-extensions"
    ]

    print(f"📋 启动命令: {' '.join(cmd)}")

    if os.name == 'nt':
        proc = subprocess.Popen(
            cmd,
            creationflags=subprocess.CREATE_NEW_CONSOLE
        )
        print(f"🖥️  启动 Chromium: CDP={cdp_port}, PID={proc.pid}")
    else:
        proc = subprocess.Popen(cmd)
        print(f"🖥️  启动 Chromium: CDP={cdp_port}, PID={proc.pid}")

    return proc


def start_mcp_instance(port, host, user_id, mcp_version, base_dir, cdp_endpoint):
    """启动单个 MCP 实例（前台保留窗口）"""

    # 使用配置文件中的DATA_PATH作为基础路径
    user_data_dir = os.path.join(DATA_PATH, 'user_data', f'user_{user_id}')
    os.makedirs(user_data_dir, exist_ok=True)

    print(f"📁 mcp使用用户数据目录: {user_data_dir}")

    width, height = get_screen_size()

    args = [
        'npx', mcp_version,
        '--port', str(port),
        '--host', host,
        '--user-data-dir', user_data_dir,
        '--isolated',
        '--viewport-size', f'{width},{height}',
        '--cdp-endpoint', cdp_endpoint,
        '--allowed-hosts', '*'
    ]

    if os.name == 'nt':
        cmd_line = subprocess.list2cmdline(args)
        print(f"🔧 启动 MCP: 端口={port} → CDP={cdp_endpoint}")
        process = subprocess.Popen([
            'cmd', '/c', 'start', '', 'cmd', '/k', cmd_line
        ])
        time.sleep(1)  # 小延迟避免窗口重叠
        return process
    else:
        print(f"🔧 启动 MCP: 端口={port} → CDP={cdp_endpoint}")
        return subprocess.Popen(args)


def main():
    """主函数，整合了浏览器启动、MCP启动和端口转发功能"""
    args = parse_arguments()
    script_dir = os.path.dirname(os.path.abspath(__file__))

    args.enable_forward = 'true'

    print("=" * 60)
    print("🔥 启动 Chromium + MCP 集群（先浏览器，后MCP）")
    print("=" * 60)
    print(f"📊 实例数量: {args.count}")
    print(f"🌐 MCP 端口: {args.mcp_start_port} ~ {args.mcp_start_port + args.count - 1}")
    print(f"🔧 CDP 端口: {args.cdp_start_port} ~ {args.cdp_start_port + args.count - 1}")
    print(f"🏠 用户ID: {args.start_user} ~ {args.start_user + args.count - 1}")
    print(f"👁️  是否启动chromium: {'是' if args.is_start_cdp else '否'}")
    print(f"👁️  是否启动mcp: {'是' if args.is_start_mcp else '否'}")
    print(f"👁️  无头模式: {'是' if args.headless else '否'}")
    print(
        f"🔄 端口转发: {'是' if args.enable_forward else '否'} (监听:{args.forward_port} → 目标:{args.forward_target_host}:{args.forward_target_port})")
    print("-" * 60)

    # 启动端口转发（如果启用）
    for i in range(args.count):
        forward_port = args.forward_port + i
        forward_target_port = args.forward_target_port + i
        if args.enable_forward:
            print("\n🔄 正在启动端口转发服务...")
            start_port_forwarding(forward_port, args.forward_target_host, forward_target_port)

    args.is_start_cdp = "true"
    args.is_start_mcp = "true"

    # =============== 阶段1：启动所有 Chromium ===============
    chromium_processes = []
    if args.is_start_cdp:
        print("\n🟢 正在启动 Chromium 浏览器...")

        for i in range(args.count):
            cdp_port = args.cdp_start_port + i
            user_id = args.start_user + i

            # 使用配置文件中的DATA_PATH作为基础路径
            user_data_dir = os.path.join(DATA_PATH, 'user_data', f'user_{user_id}')
            os.makedirs(user_data_dir, exist_ok=True)

            try:
                proc = start_chromium_instance(cdp_port, user_data_dir, headless=args.headless)
                chromium_processes.append(proc)
            except Exception as e:
                print(f"❌ 启动 Chromium 实例失败 (CDP={cdp_port}): {e}")
                continue

            time.sleep(args.delay)

        print("⏳ 等待 5 秒，确保所有浏览器完全启动...")
        time.sleep(5)

    # =============== 阶段2：启动所有 MCP ===============
    mcp_processes = []
    if args.is_start_mcp:
        print("\n🟢 正在启动 MCP 服务器...")

        # 检查 npx 是否可用
        if not check_npx_availability():
            print("⚠️  npx 不可用，但仍尝试启动 MCP 实例...")
            # 继续尝试，不直接退出

        for i in range(args.count):
            mcp_port = args.mcp_start_port + i
            cdp_port = args.cdp_start_port + i
            cdp_endpoint = f"http://{args.host}:{cdp_port}"
            user_id = args.start_user + i

            try:
                proc = start_mcp_instance(mcp_port, args.host, user_id, args.mcp_version, script_dir, cdp_endpoint)
                mcp_processes.append(proc)
            except Exception as e:
                print(f"❌ 启动 MCP 实例失败 (端口={mcp_port}): {e}")
                continue

            time.sleep(args.delay)

    # =============== 完成提示 ===============
    print("\n" + "=" * 60)
    print("🎉 启动完成！")

    if args.is_start_cdp and args.is_start_mcp:
        for i in range(args.count):
            mcp_port = args.mcp_start_port + i
            cdp_port = args.cdp_start_port + i
            print(f"🔗 MCP:{mcp_port} → CDP:localhost:{cdp_port}")

    if args.enable_forward:
        print(f"🔄 端口转发: *:{args.forward_port} → {args.forward_target_host}:{args.forward_target_port}")

    print("=" * 60)
    print("🔁 所有服务正在运行中... 按 Ctrl+C 停止服务")
    print("💡 提示：Chromium 和 MCP 窗口需手动关闭")

    try:
        # 阻塞主线程，防止 daemon 线程被杀死
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n🛑 正在退出...")
        # 可选：终止子进程（需保存 proc 列表）
        sys.exit(0)


if __name__ == "__main__":
    main()